import os
import torch
from torch.utils.data import DataLoader
import pandas as pd
from tqdm.auto import tqdm
from model_define import FoodCNN
from train_transform import FoodDataset, test_tfm
from model_vgg import ModularVGG11

def predict(model, test_loader, device):
    model.eval()
    predictions = []
    filenames = []
    with torch.no_grad():
        for images, names in tqdm(test_loader, desc='预测中'):
            images = images.to(device)
            outputs = model(images)
            _, predicted = torch.max(outputs, 1)
            predictions.extend(predicted.cpu().numpy())
            filenames.extend(names)

    # 创建结果
    results = []
    for i, (fn, pred) in enumerate(zip(filenames, predictions)):
        img_id = os.path.splitext(fn)[0]  #
        results.append({'Id': img_id, 'Category': int(pred)})

    result_df = pd.DataFrame(results)
    result_df.to_csv('ans_ours_vgg.csv', index=False)
    print("预测结果已保存到ans_ours.csv")
    return result_df
# 主执行函数
def main():
    # 加载最佳模型进行预测
    print("加载最佳模型进行预测...")
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    best_model = ModularVGG11(num_classes=11)
    best_model.load_state_dict(torch.load('best_food_model_vgg.pth'))
    best_model.to(device)
    test_dataset = FoodDataset('food11/versions/1/evaluation', 'test', test_tfm)
    test_loader = DataLoader(test_dataset, batch_size=128, shuffle=False, num_workers=2)

    # 进行预测
    predictions = predict(best_model, test_loader, device)
    print("预测完成！")
    print(predictions)

if __name__ == "__main__":
    main()